import numpy as np
import matplotlib.pyplot as plt
from scipy.integrate import odeint
from scipy.optimize import curve_fit

plt.rcParams['font.sans-serif'] = ['SimHei']  # 用黑体显示中文
plt.rcParams['axes.unicode_minus'] = False    # 正常显示负号

# 物理常数
g = 9.8  # 重力加速度 (m/s^2)

# 系统参数
l = 1.0    # 摆长 (m)
m = 0.1    # 质量 (kg)

# 有阻尼单摆的微分方程
def pendulum_eq(y, t, k):
    """
    有阻尼单摆的微分方程
    y[0] = theta (角度)
    y[1] = dtheta/dt (角速度)
    k = 阻尼系数
    """
    theta, omega = y
    dtheta_dt = omega
    domega_dt = -(k/m) * omega - (g/l) * np.sin(theta)
    return [dtheta_dt, domega_dt]

# 生成模拟实验数据
def generate_simulated_data(k_true, theta0, omega0, t_max, noise_level=0.05):
    """
    生成模拟实验数据
    k_true: 真实的阻尼系数
    theta0, omega0: 初始条件
    t_max: 时间长度
    noise_level: 噪声水平
    """
    # 时间点
    t = np.linspace(0, t_max, 200)
    
    # 真实解
    y0 = [theta0, omega0]
    sol_true = odeint(pendulum_eq, y0, t, args=(k_true,))
    theta_true = sol_true[:, 0]
    
    # 添加噪声模拟实验误差
    noise = np.random.normal(0, noise_level, len(t))
    theta_noisy = theta_true + noise
    
    return t, theta_noisy, theta_true

# 拟合函数
def fit_model(t, theta_data, theta0, omega0):
    """
    拟合阻尼系数k
    """
    def model(t, k):
        y0 = [theta0, omega0]
        sol = odeint(pendulum_eq, y0, t, args=(k,))
        return sol[:, 0]
    
    # 使用curve_fit进行参数拟合
    popt, pcov = curve_fit(model, t, theta_data, p0=[0.1])
    k_fitted = popt[0]
    k_error = np.sqrt(np.diag(pcov))[0]
    
    return k_fitted, k_error

# 绘制相图
def plot_phase_diagram(k, theta0, omega0, t_max=20):
    """
    绘制相图（角度-角速度图）
    """
    t = np.linspace(0, t_max, 1000)
    y0 = [theta0, omega0]
    sol = odeint(pendulum_eq, y0, t, args=(k,))
    theta = sol[:, 0]
    omega = sol[:, 1]
    
    plt.figure(figsize=(10, 8))
    plt.plot(theta, omega, 'b-', linewidth=2)
    plt.xlabel('角度 θ (弧度)', fontsize=14)
    plt.ylabel('角速度 dθ/dt (弧度/秒)', fontsize=14)
    plt.title(f'有阻尼单摆相图 (k = {k:.3f} kg/s)', fontsize=16)
    plt.grid(True, alpha=0.3)
    plt.xlim(-2*np.pi, 2*np.pi)
    
    # 标记起始点
    plt.plot(theta0, omega0, 'ro', markersize=10, label='起始点')
    plt.legend()
    
    # 添加平衡点
    equilibrium_points = [0, np.pi, -np.pi]
    for point in equilibrium_points:
        plt.plot(point, 0, 'ks', markersize=8)
    
    plt.text(0.2, 0.2, '稳定平衡点', fontsize=12)
    plt.text(np.pi+0.2, 0.2, '不稳定平衡点', fontsize=12)
    plt.text(-np.pi+0.2, 0.2, '不稳定平衡点', fontsize=12)
    
    plt.tight_layout()
    # 在当前目录保存图片
    plt.savefig('damped_pendulum_phase.png', dpi=150, bbox_inches='tight')
    plt.show()

# 创建单摆动画
def create_animation(k, theta0, omega0, t_max=10):
    """
    创建单摆运动动画
    """
    import matplotlib.animation as animation
    
    # 计算轨迹
    t = np.linspace(0, t_max, 500)
    y0 = [theta0, omega0]
    sol = odeint(pendulum_eq, y0, t, args=(k,))
    theta = sol[:, 0]
    
    # 设置图形
    fig, ax = plt.subplots(figsize=(8, 8))
    ax.set_xlim(-1.5, 1.5)
    ax.set_ylim(-1.5, 1.5)
    ax.set_aspect('equal')
    ax.grid(True, alpha=0.3)
    ax.set_xlabel('x (m)', fontsize=12)
    ax.set_ylabel('y (m)', fontsize=12)
    ax.set_title('有阻尼单摆运动动画', fontsize=14)
    
    # 固定点
    fixed_point, = ax.plot(0, 0, 'ko', markersize=10)
    
    # 摆杆
    rod, = ax.plot([], [], 'o-', lw=2, markersize=10)
    
    # 轨迹
    trajectory, = ax.plot([], [], 'r-', lw=1, alpha=0.5)
    
    # 轨迹存储
    x_trace, y_trace = [], []
    
    def animate(i):
        # 计算摆的位置
        x = l * np.sin(theta[i])
        y = -l * np.cos(theta[i])
        
        # 更新摆杆
        rod.set_data([0, x], [0, y])
        
        # 更新轨迹
        x_trace.append(x)
        y_trace.append(y)
        if len(x_trace) > 200:  # 只保留最近的200个点
            x_trace.pop(0)
            y_trace.pop(0)
        trajectory.set_data(x_trace, y_trace)
        
        return rod, trajectory
    
    # 创建动画
    anim = animation.FuncAnimation(fig, animate, frames=len(t), interval=50, blit=True, repeat=True)
    
    # 保存动画
    anim.save('damped_pendulum_animation.gif', writer='pillow', fps=20)
    plt.show()
    
    return anim

# 主函数
def main():
    print("有空气阻力的单摆运动仿真")
    print("=" * 40)
    
    # 真实参数
    k_true = 0.2  # 真实阻尼系数 (kg/s)
    theta0 = 0.3  # 初始角度 (弧度)
    omega0 = 0.0  # 初始角速度 (弧度/秒)
    
    print(f"真实阻尼系数: k = {k_true} kg/s")
    print(f"初始条件: θ₀ = {theta0:.2f} rad, ω₀ = {omega0:.2f} rad/s")
    
    # 1. 生成模拟实验数据
    print("\n1. 生成模拟实验数据...")
    t, theta_noisy, theta_true = generate_simulated_data(k_true, theta0, omega0, t_max=20, noise_level=0.05)
    
    # 绘制实验数据
    plt.figure(figsize=(12, 6))
    plt.plot(t, theta_noisy, 'o', markersize=4, alpha=0.7, label='实验数据')
    plt.plot(t, theta_true, '-', linewidth=2, label='真实值')
    plt.xlabel('时间 t (秒)', fontsize=14)
    plt.ylabel('角度 θ (弧度)', fontsize=14)
    plt.title('模拟实验数据', fontsize=16)
    plt.legend()
    plt.grid(True, alpha=0.3)
    plt.tight_layout()
    plt.savefig('simulated_data.png', dpi=150, bbox_inches='tight')
    plt.show()
    
    # 2. 参数拟合
    print("\n2. 参数拟合...")
    k_fitted, k_error = fit_model(t, theta_noisy, theta0, omega0)
    print(f"拟合阻尼系数: k = {k_fitted:.4f} ± {k_error:.4f} kg/s")
    print(f"相对误差: {abs(k_fitted - k_true) / k_true * 100:.2f}%")
    
    # 3. 绘制相图
    print("\n3. 绘制相图...")
    plot_phase_diagram(k_fitted, theta0, omega0, t_max=20)
    
    # 4. 创建动画
    print("\n4. 创建动画...")
    print("正在生成动画，请稍候...")
    anim = create_animation(k_fitted, theta0, omega0, t_max=10)
    print("动画已保存为 '有阻尼单摆动画.gif'")
    
    print("\n完成！")

if __name__ == "__main__":
    main()